#include <stdio.h>
#include <stdlib.h>
#include <string.h>
#include <direct.h> // Windows 目录操作
#include <io.h>     // Windows 文件操作
#include <sys/stat.h>

#define MAX_DICT_SIZE 4096
#define MAX_CODE_LEN 12
#define FILE_HEADER_MAGIC "RLEZIP"
#define FILE_HEADER_VERSION 1
#define MAX_PATH_LENGTH 4096

typedef unsigned short code_t;

typedef struct {
    char *sequence;
    int length;
} DictionaryEntry;

DictionaryEntry dictionary[MAX_DICT_SIZE];
int next_code = 256;

// 初始化字典表
void init_dictionary() {
	int i=0;
    for ( i = 0; i < 256; i++) {
        dictionary[i].sequence = (char*)malloc(1);
        dictionary[i].sequence[0] = (char)i;
        dictionary[i].length = 1;
    }
    for ( i = 256; i < MAX_DICT_SIZE; i++) {
        dictionary[i].sequence = NULL;
        dictionary[i].length = 0;
    }
}

// 释放字典表内存
void free_dictionary() {
	int i=0;
    for ( i = 0; i < MAX_DICT_SIZE; i++) {
        if (dictionary[i].sequence != NULL) {
            free(dictionary[i].sequence);
        }
    }
}

// 从输入文件中读取一个编码
code_t read_code(FILE *input, int *bit_offset, unsigned char *buffer) {
    code_t code = 0;
    int i=0;
    for ( i = 0; i < MAX_CODE_LEN; i++) {
        if (*bit_offset == 0) {
            if (fread(buffer, 1, 1, input) != 1) {
                return MAX_DICT_SIZE - 1; // 文件结束标记
            }
        }
        if (*buffer & (1 << *bit_offset)) {
            code |= (1 << i);
        }
        (*bit_offset)++;
        if (*bit_offset == 8) {
            *bit_offset = 0;
        }
    }
    return code;
}

// 解压文件数据
void decompress_file(FILE *input, FILE *output) {
    init_dictionary();
    int bit_offset = 0;
    unsigned char buffer = 0;
    code_t prev_code, curr_code;

    prev_code = read_code(input, &bit_offset, &buffer);
    if (prev_code >= next_code) {
        fprintf(stderr, "Invalid compressed data\n");
        return;
    }
    fwrite(dictionary[prev_code].sequence, 1, dictionary[prev_code].length, output);

    while ((curr_code = read_code(input, &bit_offset, &buffer)) != MAX_DICT_SIZE - 1) {
        if (curr_code < next_code) {
            fwrite(dictionary[curr_code].sequence, 1, dictionary[curr_code].length, output);
            if (next_code < MAX_DICT_SIZE) {
                dictionary[next_code].sequence = (char*)malloc(dictionary[prev_code].length + 1);
                memcpy(dictionary[next_code].sequence, dictionary[prev_code].sequence, dictionary[prev_code].length);
                dictionary[next_code].sequence[dictionary[prev_code].length] = dictionary[curr_code].sequence[0];
                dictionary[next_code].length = dictionary[prev_code].length + 1;
                next_code++;
            }
        } else {
            // 处理 LZW 特殊情况
            if (next_code < MAX_DICT_SIZE) {
                dictionary[next_code].sequence = (char*)malloc(dictionary[prev_code].length + 1);
                memcpy(dictionary[next_code].sequence, dictionary[prev_code].sequence, dictionary[prev_code].length);
                dictionary[next_code].sequence[dictionary[prev_code].length] = dictionary[prev_code].sequence[0];
                dictionary[next_code].length = dictionary[prev_code].length + 1;
                fwrite(dictionary[next_code].sequence, 1, dictionary[next_code].length, output);
                next_code++;
            }
        }
        prev_code = curr_code;
    }
    free_dictionary();
}

// 创建多级目录（Windows）
void create_directories(const char *path) {
    char dir_path[MAX_PATH_LENGTH];
    strcpy(dir_path, path);
    
    // 替换路径分隔符为 Windows 格式
    char *ch;
    for ( ch = dir_path; *ch; ch++) {
        if (*ch == '/') *ch = '\\';
    }

    // 逐级创建目录
    char *dir_part = dir_path;
    while (*dir_part) {
        if (*dir_part == '\\') {
            *dir_part = '\0';
            if (access(dir_path, 0) != 0) {
                mkdir(dir_path); // Windows 的 mkdir 无需权限参数
            }
            *dir_part = '\\';
        }
        dir_part++;
    }
}

int main(int argc, char *argv[]) {
    if (argc != 2) {
        fprintf(stderr, "Usage: %s <input.rzp>\n", argv[0]);
        return EXIT_FAILURE;
    }

    FILE *input = fopen(argv[1], "rb");
    if (!input) {
        perror("Failed to open input file");
        return EXIT_FAILURE;
    }

    // 验证文件头
    char magic[6];
    fread(magic, 1, 6, input);
    if (memcmp(magic, FILE_HEADER_MAGIC, 6) != 0) {
        fprintf(stderr, "Invalid file format (magic number mismatch)\n");
        fclose(input);
        return EXIT_FAILURE;
    }

    // 验证版本号
    unsigned char version;
    fread(&version, 1, 1, input);
    if (version != FILE_HEADER_VERSION) {
        fprintf(stderr, "Unsupported version\n");
        fclose(input);
        return EXIT_FAILURE;
    }

    // 逐个解压文件
    while (1) {
        unsigned short filename_len;
        if (fread(&filename_len, sizeof(filename_len), 1, input) != 1) {
            break; // 文件结束
        }

        char filename[MAX_PATH_LENGTH];
        fread(filename, 1, filename_len, input);
        filename[filename_len] = '\0';

        long file_size;
        fread(&file_size, sizeof(file_size), 1, input);

        // 创建目录（如果需要）
        char *last_slash = strrchr(filename, '/');
        if (last_slash != NULL) {
            *last_slash = '\0';
            create_directories(filename);
            *last_slash = '/';
        }

        printf("Extracting: %s\n", filename);
        FILE *output = fopen(filename, "wb");
        if (!output) {
            perror("Failed to create output file");
            continue;
        }

        decompress_file(input, output);
        fclose(output);
    }

    fclose(input);
    printf("Decompression completed successfully.\n");
    return EXIT_SUCCESS;
}